from glob import glob
from os.path import join as pathJoin
from PIL import Image
import numpy as np
from tqdm import tqdm
import os
from functools import reduce

try:
    from AreaDetect.libs.tools import formatImage
except ImportError as e:
    from libs.tools import formatImage

def getDatas(path, format_size=(128, 128)):
    data_path_list = glob(os.path.join(path, '*'))
    datas = []
    for data_path in data_path_list:
        img = Image.open(data_path)
        if img.mode != 'RGB': img = img.convert('RGB')
        arr = formatImage(img, format_size)
        datas.append(arr)
    return datas

def datas2X(datas, input_shape=(128, 128, 3)):
    train_X = np.array(datas, dtype='float32')
    train_X = train_X.reshape((-1, ) + input_shape)
    train_X = 1 - train_X / 255
    return train_X

def getTargets(path, pcolors, format_size=(128, 128), limit=100):
    target_path_list = glob(os.path.join(path, '*'))
    targets = []
    for target_path in tqdm(target_path_list, ncols=80):
        img = Image.open(target_path)
        if img.mode != 'RGB': img = img.convert('RGB')
        # img_arr = np.array(img.resize(format_size))
        img_arr = formatImage(img, format_size)
        filterImage = lambda color: np.sum((img_arr - color)**2, axis=2)**0.5 < limit
        classes_num = len(pcolors)
        p_mask_shape = format_size + (classes_num,)
        # p_mask = reduce(lambda a, b: a | b, map(filterImage, pcolors))
        p_mask = np.stack([filterImage(pcolor) for pcolor in pcolors], axis=2)
        # n_mask = reduce(lambda a, b: a | b, map(filterImage, ncolors))
        # mask = np.zeros(p_mask.shape + (2,), dtype='uint8')
        # mask[:,:, 0] = p_mask
        # mask[:,:, 1] = n_mask
        mask = p_mask.astype('uint8')
        targets.append(mask)
    return targets, mask.shape[-1]

def targets2TrainY(targets, L=1, output_shape=(128, 128, 2)):
    # train_Y = np.array(targets, dtype='int32').reshape(-1, 128, 128, 2)
    train_Y = np.array(targets, dtype='float32').reshape(len(targets), 128, 128, -1)
    # if only_content:
    #     train_Y = train_Y[:, :, :, 0].reshape(-1, 128, 128, 1)
    return train_Y
    # D = output_shape[0] * output_shape[1]
    # K = 1 / np.sum(train_Y / train_Y.shape[0]) * D
    # train_Y = train_Y * K * L
    # return train_Y

'''
def formatImage(image, input_shape):
    src_size = np.array(image.size)
    dst_size = np.array([input_shape[0], input_shape[1]])

    if image.mode != 'RGB': image = image.convert('RGB')
    trans_array = np.zeros(input_shape, dtype='uint8') + 255

    k = dst_size / src_size.max()
    scale_size = np.around((k * src_size)).astype('int')
    trans_vect = ((dst_size - scale_size) / 2).astype('int')

    scale_image = image.resize(tuple(scale_size), resample=Image.BOX)
    scale_array = np.array(scale_image, dtype='uint8')
    trans_array[trans_vect[1]: trans_vect[1]+scale_size[1],
                trans_vect[0]: trans_vect[0]+scale_size[0]] =\
                np.array(scale_image, dtype='uint8')

    return trans_array
'''
